#!/usr/bin/env python
# -*- coding: utf-8 -*-

from PyQt5.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, QPushButton, 
                            QListWidget, QListWidgetItem, QLabel, QComboBox,
                            QProgressBar, QGroupBox, QFormLayout, QLineEdit,
                            QTextBrowser, QSplitter, QTreeWidget, QTreeWidgetItem,
                            QHeaderView, QSpacerItem, QSizePolicy, QApplication,
                            QMessageBox, QProgressDialog, QDialog, QTableWidget, QTableWidgetItem,
                            QFrame)
from PyQt5.QtCore import Qt, QSize, pyqtSignal, QThread, QObject, QTimer
from PyQt5.QtGui import QIcon, QFont
import os
import sys
from pathlib import Path
import yaml
import shutil

# 导入Ultralytics相关模块
try:
    from ultralytics import YOLO
    from ultralytics.utils.downloads import download
    ULTRALYTICS_AVAILABLE = True
except ImportError:
    ULTRALYTICS_AVAILABLE = False
    print("警告: 未找到ultralytics包，请安装: pip install ultralytics")

class ModelDownloader(QObject):
    """模型下载器线程类"""
    progress_updated = pyqtSignal(int)
    download_completed = pyqtSignal(bool, str)  # 成功状态，消息
    
    def __init__(self, model_name):
        super().__init__()
        self.model_name = model_name
        
    def run(self):
        """运行下载任务"""
        try:
            # 使用Ultralytics的下载功能
            if not ULTRALYTICS_AVAILABLE:
                self.download_completed.emit(False, "未找到ultralytics包，请安装: pip install ultralytics")
                return
            
            # 根据模型名称确定下载路径
            model_map = {
                "YOLOv10n (Nano)": "yolov10n",
                "YOLOv10s (Small)": "yolov10s",
                "YOLOv10m (Medium)": "yolov10m",
                "YOLOv10b (Base)": "yolov10b",
                "YOLOv10l (Large)": "yolov10l",
                "YOLOv10x (XLarge)": "yolov10x"
            }
            
            # 提取纯模型名称
            base_model_name = self.model_name.split(" ")[0]
            model_type = model_map.get(self.model_name, "yolov10n")  # 默认使用nano版本
            
            # 确保模型目录存在
            models_dir = os.path.join(os.getcwd(), "models")
            os.makedirs(models_dir, exist_ok=True)
            
            # 目标文件路径
            target_path = os.path.join(models_dir, f"{model_type}.pt")
            
            # 使用Hugging Face模型地址
            source_url = f"https://github.com/THU-MIG/yolov10/releases/download/v1.1/{model_type}.pt"
            
            # 直接使用urllib进行下载，以便显示进度
            from urllib.request import urlretrieve
            
            # 进度回调函数
            def report_progress(block_num, block_size, total_size):
                if total_size > 0:
                    percent = min(100, int(block_num * block_size * 100 / total_size))
                    self.progress_updated.emit(percent)
                else:
                    self.progress_updated.emit(50)  # 如果无法获取总大小，显示50%
            
            # 下载模型
            urlretrieve(source_url, target_path, report_progress)
            
            # 完成下载
            self.progress_updated.emit(100)
            self.download_completed.emit(True, f"{self.model_name} 下载完成")
            
        except Exception as e:
            import traceback
            self.download_completed.emit(False, f"下载失败: {str(e)}\n{traceback.format_exc()}")


class ModelManagerPanel(QWidget):
    """模型管理面板"""
    
    # 自定义信号
    model_downloaded = pyqtSignal(str)  # 模型下载完成信号
    
    def __init__(self):
        super().__init__()
        self.models_dir = os.path.join(os.getcwd(), "models")
        os.makedirs(self.models_dir, exist_ok=True)
        self.downloader_thread = None
        self.init_ui()
        
    def init_ui(self):
        """初始化UI"""
        # 创建主布局
        main_layout = QVBoxLayout(self)
        
        # 主内容区域
        content_widget = QWidget()
        content_layout = QHBoxLayout(content_widget)
        
        # 创建左侧模型库面板
        left_panel = QWidget()
        left_layout = QVBoxLayout(left_panel)
        
        # 标题和说明
        title_label = QLabel("YOLOv10 模型管理")
        title_label.setStyleSheet("font-size: 16px; font-weight: bold;")
        left_layout.addWidget(title_label)
        
        info_label = QLabel("通过Ultralytics库下载和管理YOLOv10模型")
        info_label.setWordWrap(True)
        left_layout.addWidget(info_label)
        
        left_layout.addSpacing(10)
        
        # 可用模型列表
        models_group = QGroupBox("可用模型")
        models_layout = QVBoxLayout(models_group)
        
        self.models_tree = QTreeWidget()
        self.models_tree.setHeaderLabels(["模型", "精度 (AP)", "状态"])
        self.models_tree.header().setSectionResizeMode(0, QHeaderView.Stretch)
        self.models_tree.setMinimumWidth(300)
        
        self.populate_models_tree()
        
        models_layout.addWidget(self.models_tree)
        
        left_layout.addWidget(models_group)
        
        # 下载按钮
        download_layout = QHBoxLayout()
        self.download_btn = QPushButton("下载选中模型")
        self.download_btn.clicked.connect(self.download_selected_model)
        download_layout.addStretch()
        download_layout.addWidget(self.download_btn)
        
        left_layout.addLayout(download_layout)
        
        # 下载进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setVisible(False)
        left_layout.addWidget(self.progress_bar)
        
        # 创建右侧模型详情面板
        right_panel = QWidget()
        right_layout = QVBoxLayout(right_panel)
        
        # 模型详情
        details_group = QGroupBox("模型详情")
        details_layout = QVBoxLayout(details_group)
        
        # 模型基本信息表单
        form_layout = QFormLayout()
        
        self.model_name_label = QLabel("未选择")
        self.model_task_label = QLabel("目标检测")
        self.model_params_label = QLabel("未选择")
        self.model_flops_label = QLabel("未选择")
        self.model_latency_label = QLabel("未选择")
        self.model_status_label = QLabel("未选择")
        
        form_layout.addRow("名称:", self.model_name_label)
        form_layout.addRow("任务:", self.model_task_label)
        form_layout.addRow("参数数量:", self.model_params_label)
        form_layout.addRow("计算量:", self.model_flops_label)
        form_layout.addRow("推理延迟:", self.model_latency_label)
        form_layout.addRow("状态:", self.model_status_label)
        
        details_layout.addLayout(form_layout)
        
        # 模型描述
        description_label = QLabel("描述:")
        self.description_browser = QTextBrowser()
        details_layout.addWidget(description_label)
        details_layout.addWidget(self.description_browser)
        
        right_layout.addWidget(details_group)
        
        # 模型操作按钮
        actions_layout = QHBoxLayout()
        
        self.export_btn = QPushButton("导出模型")
        self.export_btn.setEnabled(False)
        self.export_btn.clicked.connect(self.export_model)
        
        self.delete_btn = QPushButton("删除模型")
        self.delete_btn.setEnabled(False)
        self.delete_btn.clicked.connect(self.delete_model)
        
        self.test_btn = QPushButton("测试模型")
        self.test_btn.setEnabled(False)
        self.test_btn.clicked.connect(self.test_model)
        
        actions_layout.addWidget(self.test_btn)
        actions_layout.addWidget(self.export_btn)
        actions_layout.addWidget(self.delete_btn)
        
        right_layout.addLayout(actions_layout)
        
        # 为了填充底部空间
        right_layout.addStretch()
        
        # Ultralytics状态信息
        if ULTRALYTICS_AVAILABLE:
            status_msg = "Ultralytics库已加载，可以使用"
            status_color = "green"
        else:
            status_msg = "警告: Ultralytics库未找到，部分功能可能不可用"
            status_color = "red"
            
        status_label = QLabel(f"状态: {status_msg}")
        status_label.setStyleSheet(f"color: {status_color}; font-weight: bold;")
        right_layout.addWidget(status_label)
        
        # 添加左右面板到主布局中，分割
        splitter = QSplitter(Qt.Horizontal)
        splitter.addWidget(left_panel)
        splitter.addWidget(right_panel)
        splitter.setSizes([400, 600])
        
        content_layout.addWidget(splitter)
        main_layout.addWidget(content_widget)
        
        # 连接信号
        self.models_tree.itemClicked.connect(self.show_model_details)
        
        # 初始化时检查是否有已下载的模型
        self.refresh_models_status()
        
    def populate_models_tree(self):
        """填充模型树"""
        self.models_tree.clear()
        
        # 展示YOLOv10模型 - 显示模型名称、精度和状态
        nano = QTreeWidgetItem(self.models_tree, ["YOLOv10n (Nano)", "38.5%", "可下载"])
        small = QTreeWidgetItem(self.models_tree, ["YOLOv10s (Small)", "46.3%", "可下载"])
        medium = QTreeWidgetItem(self.models_tree, ["YOLOv10m (Medium)", "51.1%", "可下载"])
        base = QTreeWidgetItem(self.models_tree, ["YOLOv10b (Base)", "52.5%", "可下载"])
        large = QTreeWidgetItem(self.models_tree, ["YOLOv10l (Large)", "53.2%", "可下载"])
        xlarge = QTreeWidgetItem(self.models_tree, ["YOLOv10x (XLarge)", "54.4%", "可下载"])
        
        self.models_tree.expandAll()
    
    def refresh_models_status(self):
        """刷新模型状态，检查哪些已下载"""
        # 检查模型目录下的文件
        model_files = {
            "yolov10n.pt": ("YOLOv10n (Nano)", 0),
            "yolov10s.pt": ("YOLOv10s (Small)", 1),
            "yolov10m.pt": ("YOLOv10m (Medium)", 2),
            "yolov10b.pt": ("YOLOv10b (Base)", 3),
            "yolov10l.pt": ("YOLOv10l (Large)", 4),
            "yolov10x.pt": ("YOLOv10x (XLarge)", 5),
        }
        
        for file_name, (model_name, index) in model_files.items():
            file_path = os.path.join(self.models_dir, file_name)
            if os.path.exists(file_path):
                # 更新树项目
                item = self.models_tree.topLevelItem(index)
                if item:
                    item.setText(2, "已下载")
    
    def show_model_details(self, item, column):
        """显示所选模型的详细信息"""
        # 获取选中的模型名称
        model_name = item.text(0)
        model_accuracy = item.text(1)
        model_status = item.text(2)
        
        # 更新UI
        self.model_name_label.setText(model_name)
        self.model_status_label.setText(model_status)
        
        # 根据模型类型更新详细信息
        model_params = {
            "YOLOv10n (Nano)": {"params": "2.3M", "flops": "6.7G", "latency": "1.84ms"},
            "YOLOv10s (Small)": {"params": "7.2M", "flops": "21.6G", "latency": "2.49ms"},
            "YOLOv10m (Medium)": {"params": "15.4M", "flops": "59.1G", "latency": "4.74ms"},
            "YOLOv10b (Base)": {"params": "19.1M", "flops": "92.0G", "latency": "5.74ms"},
            "YOLOv10l (Large)": {"params": "24.4M", "flops": "120.3G", "latency": "7.28ms"},
            "YOLOv10x (XLarge)": {"params": "29.5M", "flops": "160.4G", "latency": "10.70ms"}
        }
        
        if model_name in model_params:
            info = model_params[model_name]
            self.model_params_label.setText(info["params"])
            self.model_flops_label.setText(info["flops"])
            self.model_latency_label.setText(info["latency"])
        else:
            self.model_params_label.setText("未知")
            self.model_flops_label.setText("未知")
            self.model_latency_label.setText("未知")
        
        # 更新描述（根据模型类型）
        if "Nano" in model_name:
            description = f"""
            YOLOv10n (Nano) 是一个轻量级的目标检测模型，适用于资源受限的设备如移动设备和边缘计算设备。
            
            特点：
            - 轻量级的模型大小（{model_params["YOLOv10n (Nano)"]["params"]}参数）
            - FLOPs: {model_params["YOLOv10n (Nano)"]["flops"]}
            - 在COCO数据集上达到{model_accuracy} AP精度
            - 低延迟: {model_params["YOLOv10n (Nano)"]["latency"]}
            
            适用场景：
            - 移动设备应用
            - 嵌入式系统
            - 实时检测任务（低精度要求）
            
            权衡：速度优先，精度较低
            """
        elif "Small" in model_name:
            description = f"""
            YOLOv10s (Small) 是一个小型目标检测模型，在速度和精度之间取得良好平衡。
            
            特点：
            - 小型模型大小（{model_params["YOLOv10s (Small)"]["params"]}参数）
            - FLOPs: {model_params["YOLOv10s (Small)"]["flops"]}
            - 在COCO数据集上达到{model_accuracy} AP精度
            - 低延迟: {model_params["YOLOv10s (Small)"]["latency"]}
            
            适用场景：
            - 普通计算设备
            - 通用目标检测任务
            - 实时应用需求
            
            权衡：平衡速度和精度
            """
        elif "Medium" in model_name:
            description = f"""
            YOLOv10m (Medium) 是一个中型目标检测模型，提供更高的精度和适中的速度。
            
            特点：
            - 中等模型大小（{model_params["YOLOv10m (Medium)"]["params"]}参数）
            - FLOPs: {model_params["YOLOv10m (Medium)"]["flops"]}
            - 在COCO数据集上达到{model_accuracy} AP精度
            - 低延迟: {model_params["YOLOv10m (Medium)"]["latency"]}
            
            适用场景：
            - 需要较高精度的应用
            - 具备中等计算能力的设备
            - 可接受轻微延迟的场景
            """
        elif "Base" in model_name:
            description = f"""
            YOLOv10b (Base) 是一个基础型目标检测模型，提供良好的精度与速度平衡。
            
            特点：
            - 适中模型大小（{model_params["YOLOv10b (Base)"]["params"]}参数）
            - FLOPs: {model_params["YOLOv10b (Base)"]["flops"]} 
            - 在COCO数据集上达到{model_accuracy} AP精度
            - 低延迟: {model_params["YOLOv10b (Base)"]["latency"]}
            
            适用场景：
            - 需要较高精度的应用
            - 需要在多种设备上部署
            - 通用对象检测任务
            """
        elif "Large" in model_name:
            description = f"""
            YOLOv10l (Large) 是一个大型目标检测模型，提供高精度检测。
            
            特点：
            - 大型模型（{model_params["YOLOv10l (Large)"]["params"]}参数）
            - FLOPs: {model_params["YOLOv10l (Large)"]["flops"]}
            - 在COCO数据集上达到{model_accuracy} AP精度
            - 低延迟: {model_params["YOLOv10l (Large)"]["latency"]}
            
            适用场景：
            - 需要高精度的应用
            - 复杂场景下的目标检测
            - 具备良好计算资源的设备
            """
        elif "XLarge" in model_name:
            description = f"""
            YOLOv10x (XLarge) 是一个超大型目标检测模型，提供最高精度检测。
            
            特点：
            - 大型模型（{model_params["YOLOv10x (XLarge)"]["params"]}参数）
            - FLOPs: {model_params["YOLOv10x (XLarge)"]["flops"]}
            - 在COCO数据集上达到{model_accuracy} AP精度
            - 低延迟: {model_params["YOLOv10x (XLarge)"]["latency"]}
            
            适用场景：
            - 要求极高精度的应用
            - 需要检测复杂场景和小目标
            - 具备强大计算资源的设备
            
            权衡：精度优先，速度较慢
            """
        else:
            description = f"{model_name} 是一个高质量的目标检测模型，提供出色的检测性能。"
            
        self.description_browser.setText(description)
        
        # 启用或禁用按钮，基于模型状态
        is_downloaded = "已下载" in model_status
        self.export_btn.setEnabled(is_downloaded)
        self.delete_btn.setEnabled(is_downloaded)
        self.test_btn.setEnabled(is_downloaded)
        
    def set_status(self, message, is_error=False):
        """在状态栏显示消息"""
        # 查找父窗口的状态栏
        main_window = self.window()
        if hasattr(main_window, 'statusBar'):
            # 作为属性访问而不是方法调用
            status_bar = main_window.statusBar
            status_bar.showMessage(message, 5000)
        
    def download_selected_model(self):
        """下载选中的模型"""
        selected_items = self.models_tree.selectedItems()
        if not selected_items:
            self.set_status("请先选择要下载的模型", True)
            return
            
        selected_model = selected_items[0]
        model_name = selected_model.text(0)
        
        # 已下载就不再下载
        if "已下载" in selected_model.text(2):
            self.set_status(f"{model_name} 已经下载")
            return
        
        # 显示进度条
        self.progress_bar.setVisible(True)
        self.progress_bar.setValue(0)
        self.download_btn.setEnabled(False)
        self.set_status(f"正在下载 {model_name}...")
        
        # 创建下载线程
        self.downloader = ModelDownloader(model_name)
        self.thread = QThread()
        self.downloader.moveToThread(self.thread)
        
        # 连接信号
        self.thread.started.connect(self.downloader.run)
        self.downloader.progress_updated.connect(self.update_progress)
        self.downloader.download_completed.connect(self.download_finished)
        
        # 启动线程
        self.thread.start()
        
    def update_progress(self, value):
        """更新进度条"""
        self.progress_bar.setValue(value)
        
    def download_finished(self, success, message):
        """下载完成处理"""
        self.thread.quit()
        self.thread.wait()
        
        self.download_btn.setEnabled(True)
        
        if success:
            selected_items = self.models_tree.selectedItems()
            if selected_items:
                selected_model = selected_items[0]
                selected_model.setText(2, "已下载")
                self.model_status_label.setText("已下载")
                
                # 启用按钮
                self.export_btn.setEnabled(True)
                self.delete_btn.setEnabled(True)
                self.test_btn.setEnabled(True)
                
                # 发出信号
                self.model_downloaded.emit(selected_model.text(0))
                
                self.set_status(message)
        else:
            self.set_status(message, True)
            
        # 隐藏进度条
        self.progress_bar.setVisible(False)
        
    def get_model_path(self, model_name):
        """根据模型名称获取文件路径"""
        model_map = {
            "YOLOv10n (Nano)": "yolov10n.pt",
            "YOLOv10s (Small)": "yolov10s.pt",
            "YOLOv10m (Medium)": "yolov10m.pt",
            "YOLOv10b (Base)": "yolov10b.pt",
            "YOLOv10l (Large)": "yolov10l.pt",
            "YOLOv10x (XLarge)": "yolov10x.pt"
        }
        
        if model_name in model_map:
            return os.path.join(self.models_dir, model_map[model_name])
        return None
        
    def delete_model(self):
        """删除所选模型"""
        selected_items = self.models_tree.selectedItems()
        if not selected_items:
            return
            
        selected_model = selected_items[0]
        model_name = selected_model.text(0)
        
        # 确认删除
        reply = QMessageBox.question(self, "确认删除", 
                                    f"确定要删除模型 {model_name} 吗？",
                                    QMessageBox.Yes | QMessageBox.No)
        
        if reply != QMessageBox.Yes:
            return
            
        # 获取模型文件路径
        model_path = self.get_model_path(model_name)
        
        if model_path and os.path.exists(model_path):
            try:
                os.remove(model_path)
                # 更新状态
                selected_model.setText(2, "可下载")
                self.model_status_label.setText("可下载")
                
                # 禁用按钮
                self.export_btn.setEnabled(False)
                self.delete_btn.setEnabled(False)
                self.test_btn.setEnabled(False)
                
                self.set_status(f"模型 {model_name} 已删除")
            except Exception as e:
                self.set_status(f"删除失败: {str(e)}", True)
        else:
            self.set_status("无法找到模型文件", True)
        
    def export_model(self):
        """导出所选模型为其他格式"""
        selected_items = self.models_tree.selectedItems()
        if not selected_items:
            return
            
        selected_model = selected_items[0]
        model_name = selected_model.text(0)
        
        # 获取模型文件路径
        model_path = self.get_model_path(model_name)
        
        if not model_path or not os.path.exists(model_path):
            QMessageBox.warning(self, "警告", "模型文件不存在")
            return
            
        # 选择导出格式
        from PyQt5.QtWidgets import QDialog, QVBoxLayout, QRadioButton, QDialogButtonBox
        
        dialog = QDialog(self)
        dialog.setWindowTitle("选择导出格式")
        layout = QVBoxLayout(dialog)
        
        formats = ["ONNX", "TorchScript", "TensorRT", "CoreML", "TFLite"]
        radio_buttons = []
        
        for fmt in formats:
            rb = QRadioButton(fmt)
            layout.addWidget(rb)
            radio_buttons.append(rb)
        
        # 默认选择第一个
        if radio_buttons:
            radio_buttons[0].setChecked(True)
        
        button_box = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
        button_box.accepted.connect(dialog.accept)
        button_box.rejected.connect(dialog.reject)
        layout.addWidget(button_box)
        
        if dialog.exec_() == QDialog.Accepted:
            selected_format = "ONNX"  # 默认
            for rb in radio_buttons:
                if rb.isChecked():
                    selected_format = rb.text()
                    break
            
            # 这里应该实现实际的导出逻辑
            # 使用Ultralytics API导出不同格式
            if ULTRALYTICS_AVAILABLE:
                try:
                    QMessageBox.information(
                        self, 
                        "导出操作", 
                        f"模型 {model_name} 将被导出为 {selected_format} 格式\n"
                        f"(此操作在实际应用中会使用YOLO.export()方法)"
                    )
                except Exception as e:
                    QMessageBox.critical(self, "导出错误", str(e))
            else:
                QMessageBox.warning(
                    self, 
                    "无法导出", 
                    "未找到Ultralytics库，无法进行导出操作"
                )
                
    def test_model(self):
        """测试所选模型"""
        selected_items = self.models_tree.selectedItems()
        if not selected_items:
            return
            
        selected_model = selected_items[0]
        model_name = selected_model.text(0)
        
        # 获取模型文件路径
        model_path = self.get_model_path(model_name)
        
        if not model_path or not os.path.exists(model_path):
            QMessageBox.warning(self, "警告", "模型文件不存在")
            return
            
        # 实际测试逻辑
        if not ULTRALYTICS_AVAILABLE:
            QMessageBox.warning(self, "警告", "未找到ultralytics库，无法进行测试")
            return
            
        # 使用QProgressDialog显示测试进度
        progress = QProgressDialog("正在测试模型...", "取消", 0, 100, self)
        progress.setWindowTitle("模型测试")
        progress.setWindowModality(Qt.WindowModal)
        progress.setAutoClose(True)
        progress.setMinimumDuration(0)
        progress.setValue(10)
        
        try:
            # 创建测试线程
            class TestThread(QThread):
                test_completed = pyqtSignal(bool, str, dict)  # 成功状态，消息，结果
                progress_updated = pyqtSignal(int)  # 进度更新信号
                
                def __init__(self, model_path):
                    super().__init__()
                    self.model_path = model_path
                    self.using_yolov10 = False  # 默认非YOLOv10
                    
                def run(self):
                    try:
                        # 发送进度更新 - 开始
                        self.progress_updated.emit(30)
                        
                        # 尝试加载模型
                        import torch
                        
                        # 创建简单的测试数据目录
                        test_dir = os.path.join(os.getcwd(), "test_data")
                        os.makedirs(test_dir, exist_ok=True)
                        
                        # 发送进度更新
                        self.progress_updated.emit(40)
                        
                        # 创建简单的测试数据YAML
                        test_yaml = os.path.join(test_dir, "test.yaml")
                        with open(test_yaml, 'w') as f:
                            f.write("""
path: ./test_data
train: images/train
val: images/val
test: images/test

nc: 80
names: [
  'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
  'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
  'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
  'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
  'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
  'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
  'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
  'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
  'hair drier', 'toothbrush'
]
                            """)
                            
                        # 发送进度更新
                        self.progress_updated.emit(50)
                        
                        # 创建一个简单的测试示例
                        from ultralytics import YOLO
                        
                        # 发送进度更新 - 加载模型
                        self.progress_updated.emit(60)
                        
                        # 确保加载完整路径的模型文件
                        model_path = os.path.abspath(self.model_path)
                        print(f"加载模型: {model_path}")
                        
                        # 直接尝试使用 YOLOv10 模型
                        print(f"尝试加载 YOLOv10 模型: {os.path.basename(model_path)}")
                        
                        # 直接尝试以常规方式加载模型
                        try:
                            model = YOLO(model_path)
                            print(f"成功加载模型: {os.path.basename(model_path)}")
                            # 标记为YOLOv10模型
                            self.using_yolov10 = True
                        except Exception as e:
                            print(f"加载 YOLOv10 模型失败: {str(e)}")
                            print("回退到使用 YOLOv8 预训练模型进行测试...")
                            model = YOLO("yolov8n.pt")
                            # 标记为YOLOv8模型
                            self.using_yolov10 = False
                        
                        # 发送进度更新 - 模型已加载
                        self.progress_updated.emit(70)
                        
                        # 只进行简单的预测测试，不实际验证（测试目的）
                        print("开始预测图像...")
                        
                        # 保存结果图片的路径
                        temp_dir = os.path.join(os.getcwd(), "temp")
                        os.makedirs(temp_dir, exist_ok=True)
                        result_image_path = os.path.join(temp_dir, "detection_result.jpg")
                        
                        # 如果结果图片已存在，先删除
                        if os.path.exists(result_image_path):
                            try:
                                os.remove(result_image_path)
                            except Exception:
                                pass
                        
                        # 预测并保存结果图片
                        result_image_dir = os.path.join(temp_dir, "predict")
                        if os.path.exists(result_image_dir):
                            import shutil
                            shutil.rmtree(result_image_dir)  # 删除之前的预测结果文件夹
                            
                        results = model.predict(
                            source="https://ultralytics.com/images/bus.jpg", 
                            conf=0.25,
                            save=True,  # 保存结果图片
                            project=temp_dir,
                            name="predict",  # 使用明确的文件夹名称
                            exist_ok=True,  # 覆盖已有文件
                        )
                        
                        # 找到生成的结果图片
                        result_image_path = None
                        predict_dir = os.path.join(temp_dir, "predict")
                        if os.path.exists(predict_dir):
                            for file in os.listdir(predict_dir):
                                if file.endswith(".jpg") or file.endswith(".png"):
                                    result_image_path = os.path.join(predict_dir, file)
                                    break
                        
                        print(f"结果图片保存至: {result_image_path}")
                        
                        # 发送进度更新 - 预测完成
                        self.progress_updated.emit(90)
                        
                        # 构建测试结果
                        test_results = {
                            "检测到的对象": len(results[0].boxes),
                            "推理时间(ms)": float(results[0].speed["inference"]),
                            "预处理时间(ms)": float(results[0].speed["preprocess"]),
                            "后处理时间(ms)": float(results[0].speed["postprocess"]),
                            "输入尺寸": f"{results[0].orig_shape[1]}x{results[0].orig_shape[0]}"
                        }
                        
                        # 尝试获取其他模型信息 - 安全地访问可能不存在的属性
                        try:
                            if hasattr(model, 'model') and hasattr(model.model, 'yaml'):
                                if isinstance(model.model.yaml, dict):
                                    if 'task' in model.model.yaml:
                                        test_results["模型任务"] = model.model.yaml["task"]
                                    if 'imgsz' in model.model.yaml:
                                        test_results["模型输入尺寸"] = model.model.yaml["imgsz"]
                        except Exception as e:
                            print(f"获取额外模型信息时出错: {str(e)}")
                        
                        # 包含结果图片路径
                        if result_image_path and os.path.exists(result_image_path):
                            test_results["_result_image_path"] = result_image_path
                            
                        # 添加模型类型标记
                        test_results["_model_type"] = "yolov10" if self.using_yolov10 else "yolov8"
                        
                        # 发送进度完成信号
                        self.progress_updated.emit(100)
                        self.test_completed.emit(True, "测试完成", test_results)
                        
                    except Exception as e:
                        import traceback
                        error_trace = traceback.format_exc()
                        print(f"测试线程错误: {str(e)}\n{error_trace}")  # 在控制台输出错误
                        self.progress_updated.emit(100)  # 确保进度条完成
                        self.test_completed.emit(False, f"测试错误: {str(e)}\n{error_trace}", {})
            
            # 创建并启动测试线程
            self.test_thread = TestThread(model_path)
            self.test_thread.test_completed.connect(lambda success, msg, results: self.show_test_results(success, msg, results, progress))
            self.test_thread.progress_updated.connect(progress.setValue)  # 连接进度更新信号
            self.test_thread.start()
            
            # 显示进度对话框
            progress.setValue(20)
            
        except Exception as e:
            progress.cancel()
            QMessageBox.critical(self, "测试错误", f"启动测试时出错: {str(e)}")
    
    def show_test_results(self, success, message, results, progress_dialog):
        """显示测试结果"""
        # 确保进度对话框显示 100%
        progress_dialog.setValue(100)
        
        # 关闭线程
        if hasattr(self, 'test_thread') and self.test_thread.isRunning():
            self.test_thread.quit()
            self.test_thread.wait()
        
        # 关闭进度对话框 - 强制关闭它，避免卡在界面上
        progress_dialog.close()
        
        if success:
            # 获取选中的模型名称
            selected_items = self.models_tree.selectedItems()
            model_name = selected_items[0].text(0) if selected_items else "未知模型"
            
            # 根据模型类型设置对话框标题
            model_type_str = ""
            if "_model_type" in results:
                if results["_model_type"] == "yolov10":
                    model_type_str = "(使用YOLOv10模型)"
                else:
                    model_type_str = "(使用YOLOv8预训练模型)"
                    
            # 创建结果对话框
            dialog = QDialog(self)
            dialog.setWindowTitle(f"模型测试结果 - {model_name} {model_type_str}")
            dialog.resize(800, 700)  # 加大对话框以便显示图片
            
            layout = QVBoxLayout(dialog)
            
            # 添加说明标签 - 根据是否使用YOLOv8显示不同信息
            if "_model_type" in results and results["_model_type"] == "yolov10":
                info_label = QLabel(
                    "成功使用 YOLOv10 模型进行测试！"
                )
                info_label.setStyleSheet("color: #15803d; background-color: #dcfce7; padding: 8px; border-radius: 4px;")
            else:
                info_label = QLabel(
                    "注意: 由于 YOLOv10 模型加载失败，这里使用 YOLOv8 预训练模型进行测试演示。"
                    "实际应用中需要针对 YOLOv10 实现特定的支持。"
                )
                info_label.setStyleSheet("color: #1d4ed8; background-color: #dbeafe; padding: 8px; border-radius: 4px;")
            
            info_label.setWordWrap(True)
            layout.addWidget(info_label)
            
            # 创建一个分割器，上面显示图片，下面显示结果数据
            splitter = QSplitter(Qt.Vertical)
            
            # 如果有结果图片，显示图片
            if "_result_image_path" in results:
                # 创建图片显示区域
                from PyQt5.QtGui import QPixmap
                
                img_label = QLabel()
                img_label.setAlignment(Qt.AlignCenter)
                
                # 加载图片
                pixmap = QPixmap(results["_result_image_path"])
                
                # 获取图片尺寸
                pixmap_width = pixmap.width()
                pixmap_height = pixmap.height()
                print(f"原始图片尺寸: {pixmap_width}x{pixmap_height}")
                
                # 保持纵横比缩放图片以适应窗口 (更大的尺寸)
                img_label.setPixmap(pixmap.scaled(600, 450, Qt.KeepAspectRatio, Qt.SmoothTransformation))
                
                # 设置最小大小以确保图片可见
                img_label.setMinimumSize(600, 400)
                
                # 创建包含图片和按钮的容器
                img_container = QWidget()
                img_container_layout = QVBoxLayout(img_container)
                img_container_layout.addWidget(img_label, 1)  # 1表示拉伸因子
                
                # 添加按钮容器
                button_container = QWidget()
                button_layout = QHBoxLayout(button_container)
                button_layout.setContentsMargins(0, 0, 0, 0)
                
                # 添加保存按钮
                save_btn = QPushButton("保存检测结果图片")
                save_btn.setFixedWidth(180)
                
                # 保存图片函数
                def save_result_image():
                    save_path, _ = QFileDialog.getSaveFileName(
                        dialog,
                        "保存检测结果图片",
                        os.path.join(os.path.expanduser("~"), "检测结果.jpg"),
                        "图片文件 (*.jpg *.png)"
                    )
                    if save_path:
                        try:
                            import shutil
                            shutil.copy2(results["_result_image_path"], save_path)
                            QMessageBox.information(dialog, "保存成功", f"检测结果图片已保存至:\n{save_path}")
                        except Exception as e:
                            QMessageBox.critical(dialog, "保存失败", f"保存图片时出错: {str(e)}")
                
                save_btn.clicked.connect(save_result_image)
                button_layout.addWidget(save_btn)
                button_layout.addStretch()
                
                # 将按钮添加到容器底部
                img_container_layout.addWidget(button_container)
                
                # 添加到分割器
                splitter.addWidget(img_container)
                
                # 从结果字典中移除图片路径，避免在表格中显示
                results.pop("_result_image_path", None)
                
            # 创建结果表格
            table_widget = QWidget()
            table_layout = QVBoxLayout(table_widget)
            
            table = QTableWidget()
            table.setColumnCount(2)
            table.setHorizontalHeaderLabels(["参数", "值"])
            table.horizontalHeader().setSectionResizeMode(0, QHeaderView.Stretch)
            table.horizontalHeader().setSectionResizeMode(1, QHeaderView.Stretch)
            
            # 添加测试结果，过滤掉以_开头的内部字段
            filtered_results = {k: v for k, v in results.items() if not k.startswith("_")}
            table.setRowCount(len(filtered_results))
            
            row = 0
            for key, value in filtered_results.items():
                table.setItem(row, 0, QTableWidgetItem(key))
                table.setItem(row, 1, QTableWidgetItem(str(value)))
                row += 1
            
            table_layout.addWidget(table)
            splitter.addWidget(table_widget)
            
            # 添加到主布局
            layout.addWidget(splitter)
            
            # 显示对话框
            dialog.exec_()
        else:
            # 创建一个更详细的错误对话框，以便查看完整的错误信息
            error_dialog = QDialog(self)
            error_dialog.setWindowTitle("测试失败")
            error_dialog.resize(600, 400)
            
            layout = QVBoxLayout(error_dialog)
            
            # 添加错误消息文本框
            error_text = QTextBrowser()
            error_text.setPlainText(message)
            error_text.setLineWrapMode(QTextBrowser.NoWrap)  # 允许水平滚动
            
            layout.addWidget(QLabel("测试模型时发生错误:"))
            layout.addWidget(error_text)
            
            # 添加关闭按钮
            close_btn = QPushButton("关闭")
            close_btn.clicked.connect(error_dialog.accept)
            layout.addWidget(close_btn)
            
            # 显示对话框
            error_dialog.exec_()


# 测试代码
if __name__ == "__main__":
    import sys
    from PyQt5.QtWidgets import QApplication
    
    app = QApplication(sys.argv)
    window = ModelManagerPanel()
    window.show()
    sys.exit(app.exec_()) 